/*
* Copyright (C) 2012-2014 qar
* License: http://www.gnu.org/licenses/gpl.html GPL version 2
*/

#include "store.h"

#include "crypto/encrypteddevice.h"
#include "outputstream.h"
#include "util.h"

bool shouldSkipCompression(QString fileName)
{
	if (fileName.endsWith(".zip", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".rar", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".gz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".bz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".bz2", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".xz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".tgz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".tbz2", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".txz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".Z", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".tlz", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".lzma", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".qar", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".7z", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".s7z", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".cab", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mp3", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".avi", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mp4", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".ogg", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".ogv", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".webm", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mov", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mp4a", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mp4v", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".jpg", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".jpeg", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".png", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".tif", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".tiff", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".cpio", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".rpm", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".srpm", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".deb", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".docx", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".pptx", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".xlsx", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".wma", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".wmv", Qt::CaseInsensitive))
		return true;
	if (fileName.endsWith(".mpg", Qt::CaseInsensitive))
		return true;
	return false;
}

// Adds a directory to the archive.
void addDir(QString dirPath, QDataStream &out, qint32 ver, QString rootPath,
            qint64 startTime, qint64 &lastProgressTime,
            qint64 &lastProgressBytes, qint64 &progressBytes,
            const QIODevice &archiveFile)
{
    QString relativeDirPath = dirPath;
    if (relativeDirPath.startsWith(rootPath)) {
        relativeDirPath = relativeDirPath.mid(rootPath.length());
    }
    displayProgress(relativeDirPath, progressBytes,
                    startTime, lastProgressTime, lastProgressBytes,
                    &archiveFile);

    // Get and write the attributes of this dir
    // Note: dirs are always written to the archive before their contents
    {
        QFileInfo dirInfo(dirPath);
        qint32 perm = dirInfo.permissions();
        QDateTime tsCreated = dirInfo.created();
        QDateTime tsModified = dirInfo.lastModified();
        QDateTime tsRead = dirInfo.lastRead();
        out << QChar('d') << relativeDirPath << perm << tsCreated << tsModified << tsRead;
        if (ver >= 2) {
            quint32 owner = dirInfo.ownerId();
            quint32 group = dirInfo.groupId();
            out << owner << group;
        }
    }
}

// Adds a symlink to the archive.
void addSymlink(QString filePath, QDataStream &out, qint32 ver, QString rootPath,
                qint64 startTime, qint64 &lastProgressTime,
                qint64 &lastProgressBytes, qint64 &progressBytes,
                const QIODevice &archiveFile)
{
    if (ver >= 2) {
        QString relativeFilePath = filePath;
        if (relativeFilePath.startsWith(rootPath)) {
            relativeFilePath = relativeFilePath.mid(rootPath.length());
        }
        displayProgress(relativeFilePath, progressBytes,
                        startTime, lastProgressTime, lastProgressBytes,
                        &archiveFile);
        QString target = getLinkTarget(filePath);
        bool relativeToRoot = false;
        if (target.startsWith(rootPath)) {
            relativeToRoot = true;
            target = target.mid(rootPath.length());
            while (target.startsWith(QDir::separator())) {
                target = target.mid(1);
            }
        }
        if (ver >= 4) {
            out << QChar('s') << relativeToRoot << relativeFilePath << target;
        } else if (ver >= 2) {
            out << QChar('s') << relativeFilePath << target;
        }
    }
}

// Adds a file to the archive.
void addFile(QFileInfo fileInfo, QDataStream &out, qint32 ver, QString rootPath,
             qint64 startTime, qint64 &lastProgressTime,
             qint64 &lastProgressBytes, qint64 &progressBytes,
             char *buffer, const int kBufSize, int compressionLevel,
             const QIODevice &archiveFile)
{
    QString filePath = fileInfo.filePath();
    QString relativeFilePath = filePath;
    if (relativeFilePath.startsWith(rootPath)) {
        relativeFilePath = relativeFilePath.mid(rootPath.length());
    }
    displayProgress(relativeFilePath, progressBytes,
                    startTime, lastProgressTime, lastProgressBytes);

    // Get file attributes and save them
    {
        qint32 perm = fileInfo.permissions();
        QDateTime tsCreated = fileInfo.created();
        QDateTime tsModified = fileInfo.lastModified();
        QDateTime tsRead = fileInfo.lastRead();
        out << QChar('f') << relativeFilePath << perm << tsCreated << tsModified << tsRead;
        if (ver >= 2) {
            quint32 owner = fileInfo.ownerId();
            quint32 group = fileInfo.groupId();
            out << owner << group;
        }
    }

    // Save the file contents
    {
        QFile file(filePath);
        if (!file.open(QIODevice::ReadOnly)) {
            qErr() << "\nWarning: Failed to open file (skipped):" << filePath;
        } else {
            const int maxThreads = 32;
            QByteArray buffers[maxThreads];
            bool isCompressed[maxThreads];
			if (compressionLevel != 0) {
				if (shouldSkipCompression(filePath)) {
					compressionLevel = 0;
				}
			}
            while (!file.atEnd()) {
                displayProgress(relativeFilePath, progressBytes,
                                startTime, lastProgressTime, lastProgressBytes,
                                &archiveFile);

                int numBuffers = 0;
                qint64 totalRead = 0;

                for (int i = 0; i < maxThreads; i++) {
                    qint64 readCount = file.read(buffer, kBufSize);
                    if (readCount <= 0)
                        break;
                    buffers[i] = QByteArray(buffer, readCount);
                    isCompressed[i] = false;
                    numBuffers++;
                    totalRead += readCount;
                }

                if (compressionLevel != 0) {
#pragma omp parallel for schedule(dynamic) default(none) shared(numBuffers, buffers, isCompressed, compressionLevel)
                    for (int i = 0; i < numBuffers; i++) {
                        QByteArray compressed = qCompress(buffers[i], compressionLevel);
                        if (compressed.count() < buffers[i].count()) {
                            buffers[i] = compressed;
                            isCompressed[i] = true;
                        }
                    }
                }

                for (int i = 0; i < numBuffers; i++) {
                    if (compressionLevel == 0) {
                        out << QChar('c') << buffers[i];
                    } else {
                        if (isCompressed[i]) {
                            out << QChar('z') << buffers[i];
                        } else {
                            out << QChar('c') << buffers[i];
                        }
                    }
                }
                progressBytes += totalRead;
            }
        }
        out << QChar('e');
    }
}

void store(QString path, int compressionLevel, bool timestamp, bool encrypt, QString encryptionKeyFile)
{
    const int kBufSize = 16 * 1024;
    char *buffer = new char[kBufSize];

    QString cleanPath = path;
    while (cleanPath.endsWith(QDir::separator()))
        cleanPath.remove(cleanPath.length() - 1, 1);
    QString archiveName = QFileInfo(cleanPath).fileName() +
                          (timestamp ? QDateTime::currentDateTime().toString("-yyyy-MM-dd-hh-mm-ss-zzz") : QString()) +
                          ".qar";
    qErr() << QString("Archiving %1 to %2 ...").arg(path).arg(archiveName);

    // Create the archive file and open it for writing
    QFile archiveFile(archiveName);
    if (archiveFile.exists()) {
        qErr() << "\nError: File already exists:" << archiveName;
        exit(-1);
    }
    if (!archiveFile.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
        qErr() << "\nError: Failed to create file:" << archiveName;
        exit(-1);
    }

    QDataStream out;
    // We always use DataStream format 4.0
    out.setVersion(QDataStream::Qt_4_0);
    out.setDevice(&archiveFile);

    // Version
    qint32 ver = 4;
    out << ver;

    qErr() << "Archive version:" << ver;
    qErr() << "Archive compression level:" << compressionLevel;

    EncryptedDevice encryptedFile(&archiveFile);
    if (ver >= 3) {
		if (!encryptedFile.open(QIODevice::WriteOnly,
								Hasher::Sha1,
								encrypt ? Cipher::Salsa20 : Cipher::NoEncryption,
								Authenticator::NoAuthentication,
								encryptionKeyFile)) {
            qErr() << "\nError: Failed to open file";
            exit(-1);
        }
        out.setDevice(&encryptedFile);
    }

    // Progress state
    qint64 startTime = currentMSecsSinceEpoch();
    qint64 lastProgressTime = startTime;
    qint64 lastProgressBytes = 0;
    qint64 progressBytes = 0;

    // The root of the directory we are archiving
    QString rootPath = QFileInfo(path).absoluteFilePath();

    // Keep a queue of all the subdirectories we need to traverse
    QList<QString> dirQueue;
    dirQueue << rootPath;

	while (!dirQueue.isEmpty()) {
        QString dirPath = dirQueue.takeFirst();
        addDir(dirPath, out, ver, rootPath,
               startTime, lastProgressTime,
               lastProgressBytes, progressBytes,
               archiveFile);

        // Open dir
        QDir dir;
        if (!dir.cd(dirPath)) {
            qErr() << "\nError: Failed to open dir:" << dirPath;
            continue;
        }

        // List subdirs
        dir.setFilter(QDir::Dirs | QDir::Hidden | QDir::NoDotAndDotDot | QDir::System);
        dir.setSorting(QDir::Size | QDir::Reversed);
        QFileInfoList list = dir.entryInfoList();
        for (int i = 0; i < list.size(); i++) {
            QFileInfo fileInfo = list.at(i);
            if (fileInfo.isSymLink()) {
                addSymlink(fileInfo.filePath(), out, ver, rootPath,
                           startTime, lastProgressTime,
                           lastProgressBytes, progressBytes,
                           archiveFile);
                // Do not follow symlinks
                continue;
            } else if (fileInfo.isDir()) {
                dirQueue << fileInfo.absoluteFilePath();
            }
        }

        // List files
        dir.setFilter(QDir::Files | QDir::Hidden | QDir::NoDotAndDotDot | QDir::System);
        dir.setSorting(QDir::Size | QDir::Reversed);
        list = dir.entryInfoList();
        for (int i = 0; i < list.size(); i++) {
            QFileInfo fileInfo = list.at(i);
            if (fileInfo.isSymLink()) {
                addSymlink(fileInfo.filePath(), out, ver, rootPath,
                           startTime, lastProgressTime,
                           lastProgressBytes, progressBytes,
                           archiveFile);
                // Do not follow symlinks
                continue;
            } else if (fileInfo.isFile()) {
                addFile(fileInfo, out, ver, rootPath,
                        startTime, lastProgressTime,
                        lastProgressBytes, progressBytes,
                        buffer, kBufSize, compressionLevel,
                        archiveFile);
            }
        }
    }

	// We add the dirs again so that the timestamps are correct at extraction
	dirQueue << rootPath;
	while (!dirQueue.isEmpty()) {
		QString dirPath = dirQueue.takeFirst();
		addDir(dirPath, out, ver, rootPath,
			   startTime, lastProgressTime,
			   lastProgressBytes, progressBytes,
			   archiveFile);

		// Open dir
		QDir dir;
		if (!dir.cd(dirPath)) {
			qErr() << "\nError: Failed to open dir:" << dirPath;
			continue;
		}

		// List subdirs
		dir.setFilter(QDir::Dirs | QDir::Hidden | QDir::NoDotAndDotDot | QDir::System);
		dir.setSorting(QDir::Size | QDir::Reversed);
		QFileInfoList list = dir.entryInfoList();
		for (int i = 0; i < list.size(); i++) {
			QFileInfo fileInfo = list.at(i);
			if (fileInfo.isDir()) {
				dirQueue << fileInfo.absoluteFilePath();
			}
		}
	}

    delete[] buffer;
    encryptedFile.close();
    archiveFile.close();
    if (out.status() != QDataStream::Ok) {
        qErr() << "\nWrite error (archive corrupt?)";
        exit(-1);
    }
    qint64 writtenBytes = archiveFile.size();
    qint64 writtenBytesApprox = writtenBytes;
    QString writtenBytesApproxUnit;
    addSizePrefix(writtenBytesApprox, writtenBytesApproxUnit);
    qErr() << QString("\nFinished creating archive %3. Compression ratio: %4% (archive size: %1 %2B).")
              .arg(writtenBytesApprox)
              .arg(writtenBytesApproxUnit)
              .arg(archiveName)
              .arg((writtenBytes * 100) / qMax(1LL, progressBytes));
}
